Support MRL (Matryoshka Representation Learning) #676

kozistr · 2025-07-04T08:49:29Z

What does this PR do?

Fixes #673

Add dimensions field to the embed request API spec.

compatible with OpenAI client
both http & grpc
add a test case for the MRL feature
remove the unused test cases in router/tests

$ ./target/release/text-embeddings-router --model-id ../Qwen3-Embedding-0.6B/ --pooling last-token --port 8080 --dtype float32 --auto-truncate --max-batch-tokens 512

from openai import OpenAI

client = OpenAI(api_key='', base_url='http://127.0.0.1:8080')

response = client.embeddings.create(
    model='Qwen3-Embedding-0.6B',
    input='test',
    dimensions=256,
)

len(response.data[0].embedding)
256

tei_results = np.asarray(
    requests.post(
        'http://127.0.0.1:8080/embed',
        data=json.dumps({'inputs': sentences[0], 'normalize': True, 'dimensions': 256}),
        headers={'Content-type': 'application/json'},
    ).json()
)
tei_results.shape
(1, 256)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@Narsil @alvarobartt

alvarobartt

LGTM, thanks for the addition! Also I like the macros for Prometheus metrics, those look cleaner IMO.

P.S. Guessing that the snapshots were removed because those were not used, right?

router/src/http/types.rs

Co-authored-by: Alvaro Bartolome <[email protected]>

kozistr · 2025-07-09T01:05:35Z

LGTM, thanks for the addition! Also I like the macros for Prometheus metrics, those look cleaner IMO.

P.S. Guessing that the snapshots were removed because those were not used, right?

oh, sorry for the missing context. you're right!

looks like these test cases (router/tests) were all commented, so I've removed the snapshots.

updated)
I've just updated the PR description too!

thanks for the pointing this :)

michaelfeil · 2025-07-15T00:58:25Z

@kozistr QQ: what happens if you set normalize=True + mrl dimensions? Otherwise PR looks good.

kozistr · 2025-07-15T01:33:48Z

@kozistr QQ: what happens if you set normalize=True + mrl dimensions? Otherwise PR looks good.

@michaelfeil Hi! MRL will always run first before the normalization! (+ normalize is True as a default value)

thanks for double-checking this :)

michaelfeil · 2025-07-22T19:03:04Z

Is there something bringing this to mainline?

Narsil · 2025-07-23T09:54:48Z

core/src/infer.rs

+        metrics::counter!("te_embed_success").increment(1);
+        metrics::histogram!("te_embed_duration").record(total_time.as_secs_f64());
+        metrics::histogram!("te_embed_tokenization_duration")
+            .record(response.metadata.tokenization.as_secs_f64());
+        metrics::histogram!("te_embed_queue_duration")
+            .record(response.metadata.queue.as_secs_f64());
+        metrics::histogram!("te_embed_inference_duration")
+            .record(response.metadata.inference.as_secs_f64());


Can you revert those changes ? They don't seem linked to the PR.

reverted 9131071

Narsil · 2025-07-23T09:55:50Z

core/src/infer.rs

        permit: OwnedSemaphorePermit,
    ) -> Result<PooledEmbeddingsInferResponse, TextEmbeddingsError> {
        let start_time = Instant::now();

        if self.is_splade() && normalize {
-            let counter = metrics::counter!("te_request_failure", "err" => "model_type");
-            counter.increment(1);
+            metrics::counter!("te_request_failure", "err" => "model_type").increment(1);


reverted 9131071

Narsil · 2025-07-23T09:56:25Z

core/src/infer.rs

            let message = "`normalize` is not available for SPLADE models".to_string();
            tracing::error!("{message}");
            return Err(TextEmbeddingsError::Backend(BackendError::Inference(
                message,
            )));
        }

+        if let Some(dimensions) = dimensions {
+            if dimensions == 0 {
+                metrics::counter!("te_request_failure", "err" => "validation").increment(1);


Shouldn't they also be smaller than the maximum embedding dimension ?

At that time, I was considering silently returning the embedding with the original size when the given dimension is larger than the expected size, like line 274.

On second thought, like you mentioned, it'd be better to raise an error explicitly when the size is larger than expected in terms of validity.

I'll add an extra validation logic to check whether the given size is larger than the size of the embedding. thanks for catching this!

added a1b1a26

Narsil

LGTM !

kozistr added 3 commits July 4, 2025 17:42

update: remove unused snapshots

96cb307

feature: support MRL functionality

37439f5

update: test case

b5c9024

kozistr changed the title ~~Support MRL (Matryoshka Representation Learning~~ Support MRL (Matryoshka Representation Learning) Jul 4, 2025

kozistr added 2 commits July 4, 2025 23:49

update: snapshots

dd87f5c

update: grpc

3336a63

kozistr marked this pull request as ready for review July 4, 2025 15:22

alvarobartt mentioned this pull request Jul 8, 2025

Qwen3 4B model embs only 1024 dims (not 2560)? #658

Open

4 tasks

alvarobartt previously approved these changes Jul 8, 2025

View reviewed changes

router/src/http/types.rs Outdated Show resolved Hide resolved

Update router/src/http/types.rs

d56c296

Co-authored-by: Alvaro Bartolome <[email protected]>

kozistr dismissed alvarobartt’s stale review via d56c296 July 9, 2025 01:02

alvarobartt previously approved these changes Jul 9, 2025

View reviewed changes

Narsil reviewed Jul 23, 2025

View reviewed changes

revert changes

9131071

kozistr dismissed alvarobartt’s stale review via 9131071 July 23, 2025 15:49

update: validate dimension size

a1b1a26

Narsil approved these changes Jul 23, 2025

View reviewed changes

Narsil merged commit 45df4fa into huggingface:main Jul 23, 2025

kozistr deleted the feature/mrl branch July 23, 2025 16:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Support MRL (Matryoshka Representation Learning) #676

Support MRL (Matryoshka Representation Learning) #676

kozistr commented Jul 4, 2025 •

edited

Loading

Uh oh!

alvarobartt left a comment

Uh oh!

Uh oh!

kozistr commented Jul 9, 2025 •

edited

Loading

Uh oh!

michaelfeil commented Jul 15, 2025

Uh oh!

kozistr commented Jul 15, 2025

Uh oh!

michaelfeil commented Jul 22, 2025

Uh oh!

Narsil Jul 23, 2025

Uh oh!

kozistr Jul 23, 2025

Uh oh!

Narsil Jul 23, 2025

Uh oh!

kozistr Jul 23, 2025

Uh oh!

Narsil Jul 23, 2025

Uh oh!

kozistr Jul 23, 2025

Uh oh!

kozistr Jul 23, 2025

Uh oh!

Narsil left a comment

Uh oh!

Uh oh!

Support MRL (Matryoshka Representation Learning) #676

Support MRL (Matryoshka Representation Learning) #676

Conversation

kozistr commented Jul 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

alvarobartt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kozistr commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

michaelfeil commented Jul 15, 2025

Uh oh!

kozistr commented Jul 15, 2025

Uh oh!

michaelfeil commented Jul 22, 2025

Uh oh!

Narsil Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

kozistr Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Narsil Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

kozistr Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Narsil Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

kozistr Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

kozistr Jul 23, 2025

Choose a reason for hiding this comment

Uh oh!

Narsil left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kozistr commented Jul 4, 2025 •

edited

Loading

kozistr commented Jul 9, 2025 •

edited

Loading